#include "net.h"
#include "memtool.h"
#include "client.h"
#include <stdlib.h>
#include <unistd.h>
#include <assert.h>
#include <atomic>
#include <infiniband/verbs.h>

std::atomic_bool pp_bucket[PP_BUCKET_NUM];
union res_union res_pool[PP_BUCKET_NUM];

static void publish_qp_info(struct client_shared_ctx *ctx ,const char *prefix, 
                      struct ibv_qp *qp, int loc_id, int rmt_id){
  struct qp_info info;
  char name[SHORT_STR_LEN];
  info.lid = ctx->port_attr.lid;
  info.qpn = qp->qp_num;
  snprintf(name, SHORT_STR_LEN, "%s-%d-%d", prefix, loc_id, rmt_id);
  if(strlen(name) >= SHORT_STR_LEN - 1)p_die("prefix too long\n");
  p_info("publish qp :%s\n", name);
  mt_set(ctx->mi, name, (char *)&info, sizeof(struct qp_info));
}
static void get_qp_info(struct mc_ctx *mi, const char *prefix, int rmt_id, int loc_id, struct qp_info *info){
  struct qp_info *qi;
  char name[SHORT_STR_LEN];
  snprintf(name, SHORT_STR_LEN, "%s-%d-%d", prefix, rmt_id, loc_id);
  if(strlen(name) >= SHORT_STR_LEN - 1)p_die("prefix too long\n");
  p_info("get qp : %s\n", name);
  while(mt_get(mi, name, (char **)&qi, NULL))usleep(1000);
  memcpy(info, qi, sizeof(struct qp_info));
  free(qi);
}

static void connect_qp(int ib_port ,struct ibv_qp *qp, struct qp_info *info){
  struct ibv_qp_attr attr;
  int res;

  attr.qp_state = IBV_QPS_RTR;
  attr.path_mtu = IBV_MTU_4096;
  attr.dest_qp_num = info->qpn;
  attr.rq_psn = PSN;
  attr.max_dest_rd_atomic = 16;  // relate to IBV_WC_REM_INV_REQ_ERR
  attr.min_rnr_timer = 12;
  attr.ah_attr.is_global = 0;
  attr.ah_attr.dlid = info->lid;
  attr.ah_attr.sl = 0;
  attr.ah_attr.src_path_bits = 0;
  attr.ah_attr.port_num = ib_port;
  res = ibv_modify_qp(qp, &attr,
              IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU |
              IBV_QP_DEST_QPN | IBV_QP_RQ_PSN |
              IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER
              );
  assert(res == 0);
  attr.qp_state = IBV_QPS_RTS;
  attr.timeout = 14;  
  attr.retry_cnt = 7;
  attr.rnr_retry = 7;
  attr.sq_psn = PSN;
  attr.max_rd_atomic = 16; 
  res = ibv_modify_qp(qp, &attr,
              IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT |
               IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN |
              IBV_QP_MAX_QP_RD_ATOMIC);
  assert(res == 0);
}

static struct ibv_qp *create_qp_init(struct ibv_pd *pd, struct ibv_cq *recv_cq, 
                            struct ibv_cq *send_cq, int ib_port){
  struct ibv_qp_init_attr c_attr;
  struct ibv_qp *qp;
  struct ibv_qp_attr i_attr;
  int res;

  memset(&c_attr, 0, sizeof c_attr);
  c_attr.send_cq = send_cq;
  c_attr.recv_cq = recv_cq;
  c_attr.qp_type = IBV_QPT_RC;
  c_attr.cap.max_send_wr = CQ_DEPTH;
  c_attr.cap.max_recv_wr = CQ_DEPTH;
  c_attr.cap.max_send_sge = MAX_SGE_SIZE;
  c_attr.cap.max_recv_sge = MAX_SGE_SIZE;
  c_attr.cap.max_inline_data = 0; 
  c_attr.sq_sig_all = 0;
  qp = ibv_create_qp(pd, &c_attr);
  assert(qp);
  memset(&i_attr,0,sizeof(struct ibv_qp_attr));
  i_attr.qp_state = IBV_QPS_INIT;
  i_attr.port_num = ib_port;
  i_attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | 
                                IBV_ACCESS_REMOTE_WRITE |
                                IBV_ACCESS_REMOTE_READ |
                                IBV_ACCESS_REMOTE_ATOMIC;
  res = ibv_modify_qp(qp, &i_attr ,IBV_QP_STATE | 
                          IBV_QP_PKEY_INDEX | IBV_QP_PORT |
                          IBV_QP_ACCESS_FLAGS);
  assert(res == 0);
  return qp;
}
void client_conn_init(struct client_shared_ctx *ctx){
  int *i_tmp;
  int server_num, i, id, res;

  //get server num
  while(mt_get(ctx->mi, "server-num", (char **)&i_tmp, NULL))sleep(1);
  server_num = *i_tmp;
  ctx->server_num = server_num;
  //get id
  while(mt_get(ctx->mi, "client-next-id", NULL, NULL))sleep(1);
  id = mt_incr(ctx->mi, "client-next-id") - 1;
  ctx->id = id;
  p_info("client id : %d\n", id);

  //get the device and alloc the pd
  {
    struct ibv_device **dev_list;
    int num_dev;
    dev_list = ibv_get_device_list(&num_dev);
    assert(dev_list > 0);
    ctx->ctx = ibv_open_device(dev_list[0]);
    assert(ctx->ctx);
    res = ibv_query_port(ctx->ctx, ctx->ib_port, &ctx->port_attr);
    assert(!res);
    ctx->pd = ibv_alloc_pd(ctx->ctx);
    assert(ctx->pd);
  }
  //create all the QP of client
  {
    ctx->client_qp = (struct ibv_qp **)malloc(server_num * sizeof(struct ibv_qp *));
    ctx->send_cq = (struct ibv_cq **)malloc(server_num * sizeof(struct ibv_cq *));
    assert(ctx->client_qp && ctx->send_cq);
    ctx->recv_cq = ibv_create_cq(ctx->ctx, CQ_DEPTH, NULL, NULL, 0);
    assert(ctx->recv_cq);
    for(i = 0; i < server_num; i++){
      ctx->send_cq[i] = ibv_create_cq(ctx->ctx, CQ_DEPTH, NULL, NULL, 0);
      assert(ctx->send_cq[i]);
      ctx->client_qp[i] = create_qp_init(ctx->pd, ctx->recv_cq, 
                            ctx->send_cq[i], ctx->ib_port);
      assert(ctx->client_qp[i]);
    }
  }
  //publish all qps
  {
    for(i = 0; i < server_num; i++){
      publish_qp_info(ctx, "user-client", ctx->client_qp[i], id, i);
    }
  }
  //get all qp info
  {
    ctx->qp_info = (struct qp_info *)malloc(server_num * sizeof(struct qp_info));
    assert(ctx->qp_info);
    for(i = 0; i < server_num; i++){
      get_qp_info(ctx->mi, "client", i, id, &ctx->qp_info[i]);
    }
  }
  //connect server qps
  {
    for(i = 0; i < server_num; i++){
      connect_qp(ctx->ib_port, ctx->client_qp[i], &ctx->qp_info[i]);
    }
  }
  //create send id
  {
    ctx->send_id = new std::queue<int>();
    ctx->send_id_lock = new std::mutex();
    ctx->send_id_lock->lock();
    for(i = 0; i < PP_BUCKET_NUM; i++){
      ctx->send_id->push(i);
    }
    ctx->send_id_lock->unlock();
  }
}

struct recv_ring32 *create_recv_ring32(struct ibv_pd *pd, int num){
  void *mem;
  struct recv_ring32 *p = (struct recv_ring32 *)malloc(sizeof(struct recv_ring32));
  assert(p);
  mem = (struct m32 *)memalign(64, num * 32);
  assert(mem);
  p->mem = (struct m32 *)mem;
  p->num = num;
  p->head = 0;
  p->tail = 0;
  p->mr = ibv_reg_mr(pd, (void *)mem, num << 7, MR_PERMISSION);
  assert(p->mr);
  p->lkey = p->mr->lkey;
  return p;
}
void post_recv(unsigned size, int qp_id, struct ibv_qp *qp, unsigned long buf, unsigned lkey){
  struct ibv_sge li = {buf, size, lkey};
  struct ibv_recv_wr *bad_wr, recv_wr;
  int ret;
  memset(&recv_wr, 0, sizeof(struct ibv_recv_wr));
  recv_wr.wr_id = qp_id;
  recv_wr.sg_list = &li;
  recv_wr.num_sge = 1;
  ret = ibv_post_recv(qp, &recv_wr, &bad_wr);
  assert(!ret);
}
void post_recv_one_ring32(struct recv_ring32 *ring, struct ibv_qp *qp, int qp_id){
  assert((ring->head + 1) % ring->num != ring->tail);
  post_recv(32, qp_id, qp, (unsigned long)&ring->mem[ring->head], ring->lkey);
  ++ring->head;
  if(ring->head >= ring->num)ring->head -= ring->num;
}
void *get_recv_ring_addr(struct recv_ring32 *ring){
  return ring->mem + ring->tail;
}
void put_recv_one_ring32(struct recv_ring32 *ring){
  ++ring->tail;
  if(ring->tail >= ring->num)ring->tail -= ring->num;
}
void *poll_routine(void *p){
  struct client_shared_ctx *sctx = (struct client_shared_ctx *)p;
  struct ibv_cq *cq;
  struct ibv_qp **qp;
  int server_num;
  int n, i, j;
  struct recv_ring32 **rr;
  struct ibv_wc wc[POLL_NUM];
  union res_union *res;

  server_num = sctx->server_num;
  cq = sctx->recv_cq;
  qp = sctx->client_qp;
  rr = (struct recv_ring32 **)malloc(server_num * sizeof(struct recv_ring32));
  for(i = 0; i < server_num; i++)rr[i] = create_recv_ring32(sctx->pd, 64);
  for(i = 0; i < server_num; i++){
    for(j = 0; j < 16; j++){//16 post recv ahead
      post_recv_one_ring32(rr[i], qp[i], i);
    }
  }
  for(;;){
    n = ibv_poll_cq(cq, POLL_NUM, wc);
    for(i = 0; i < n; i++){
      assert(wc[i].status == 0);
      //p_info("get wc, id: %d type %d \n\n", (int)wc[i].wr_id, (int)wc[i].opcode);
      res = (union res_union *)get_recv_ring_addr(rr[wc[i].wr_id]);
      res_pool[res->res.req_id] = *res;
      pp_bucket[res->res.req_id] = true;
      put_recv_one_ring32(rr[wc[i].wr_id]);
    }
  }
}

int get_send_id(struct client_shared_ctx *sctx){
  int res = -1;
  while(res == -1){
    sctx->send_id_lock->lock();
    if(!sctx->send_id->empty()){
      res = sctx->send_id->front();
      sctx->send_id->pop();
    }
    sctx->send_id_lock->unlock();
  }
  return res;
}
void put_send_id(struct client_shared_ctx *sctx,int val){
  sctx->send_id_lock->lock();
  sctx->send_id->push(val);
  sctx->send_id_lock->unlock();
}
void post_send(unsigned size, int id, struct ibv_qp *qp, unsigned long addr, unsigned lkey){
  struct ibv_sge li = {addr, size, lkey};
  struct ibv_send_wr *bad_wr, send_wr;
  int res;
  memset(&send_wr, 0, sizeof(send_wr));
  send_wr.wr_id = id;
  send_wr.sg_list = &li;
  send_wr.num_sge = 1;
  send_wr.opcode = IBV_WR_SEND;
  send_wr.send_flags = IBV_SEND_SIGNALED;
  res = ibv_post_send(qp, &send_wr, &bad_wr);
  if(res){
    p_die("send wrong, errno %d\n", errno);
  }
}

void pingpong(struct client_thread_ctx *ctx, int server_id, union req_union *req, union res_union *res){
  int id;
  struct ibv_wc wc;
  id = get_send_id(ctx->sctx);
  pp_bucket[id] = false;
  *(ctx->req) = *req;
  ctx->req->req.req_id = id;
  post_send(32, server_id, ctx->sctx->client_qp[server_id], (unsigned long)ctx->req, ctx->req_mr->lkey);
  while(!ibv_poll_cq(ctx->sctx->send_cq[server_id], 1, &wc));
  assert(wc.status == 0);
  while(pp_bucket[id] == false);
  *res = res_pool[id];
  put_send_id(ctx->sctx, id);
  return ;
}